//
//  DemoSource.swift
//  BinaryTree
//
//  Created by 郭红旗 on 2021/12/28.
//

import UIKit

class DemoSource<T:Equatable> {
    typealias comparatorClosure = (_ e1:T,_ e2:T) -> Int
    private var comparator:comparatorClosure
    
    init(comparator:@escaping comparatorClosure) {
        self.comparator = comparator
    }

    var root:BinaryTN<T>?
    
    func printFunc() {
        print(root!.asRBString)
    }
    
    func remove(num:T) {
        guard var node = nodeForNum(num: num) else {
            return
        }
        if hasTwoChild(node: node) {
            let replace = forerunner(node: node)
            node.element = replace.element
            node = replace
        }
        
        if node.left != nil || node.right != nil {
            let leftNode = node.left
            leftNode?.parent = node.parent
            if node.isLeft {
                node.parent?.left = leftNode
            }else if node.isRight{
                node.parent?.right = leftNode
            }
            afterRemove(node: node)
        }else if node.parent == nil{
            root = nil
            afterRemove(node: nil)
        }else{
            if node.isLeft {
                node.parent?.left = nil
            }else if node.isRight{
                node.parent?.right = nil
            }
            afterRemove(node: node)
        }
    }
    
    func afterRemove(node:BinaryTN<T>?) {
        var tree:BinaryTN<T>? = node
        while tree != nil {
            if isBalance(node: tree!) {
                updateHigh(node: tree!)
            }else{
                balance(node: tree!)
            }
            tree = tree!.parent
        }
    }
    
    func hasTwoChild(node:BinaryTN<T>) -> Bool {
        if node.left != nil && node.right != nil {
            return true
        }
        return false
    }
    
    func forerunner(node:BinaryTN<T>) -> BinaryTN<T> {
        var tree = node.left
        while tree?.right != nil {
            tree = tree?.right
        }
        return tree!
    }
    
    func nodeForNum(num:T) -> BinaryTN<T>? {
        var node = root
        while node != nil {
            if node?.element == num {
                return node
            }
            let result = comparator(num,(node?.element)!)
            if result > 0 {
                node = node?.right
            }else{
                node = node?.left
            }
        }
        return nil
    }
    
    func add(num:T) {
        guard let newNode = root else {
            root = BinaryTN.init(value:num , parent: nil)
            afterAdd(node: root!)
            return
        }
        var tree:BinaryTN<T>? = newNode
        var result:Int = 0
        var replace:BinaryTN<T>? = newNode
        while tree != nil {
            result = comparator(num,tree!.element!)
            replace = tree
            if result < 0 {
                tree = tree!.left
            }else if result > 0 {
                tree = tree!.right
            }
        }
        let newTree = BinaryTN.init(value: num, parent: replace)
        if result < 0 {
            replace?.left = newTree
        }else if result > 0{
            replace?.right = newTree
        }
        afterAdd(node: newTree)
    }
    func afterAdd(node:BinaryTN<T>?) {
        guard let tree = node?.parent else {
            root?.blackFun()
            return
        }
        if tree.isBlack() { return }
        /// 父节点为红色
        if let sibing = tree.sibing() {
            if sibing.isRed() {
                sibing.blackFun()
                tree.blackFun()
                tree.parent?.redFun()
                afterAdd(node: tree.parent)
                return
            }
        }
        
        let grand = tree.parent
        if tree.element == grand?.left?.element { // L
            if node?.element == tree.left?.element { // LL
                tree.blackFun()
                grand?.redFun()
                rotateRight(node: grand!)
            }else{ // LR
                grand?.redFun()
                node?.blackFun()
                rotateLeft(node: tree)
                rotateRight(node: grand!)
            }
        }else{
            if node?.element == tree.right?.element { // RR
                tree.blackFun()
                grand?.redFun()
                rotateLeft(node: grand!)
            }else{ // RL
                grand?.redFun()
                node?.blackFun()
                rotateRight(node: tree)
                rotateLeft(node: grand!)
            }
        }
        
        
    }
//    func afterAdd(node:BinaryTN<T>) {
//        var tree:BinaryTN<T>? = node
//        while tree != nil {
//            if isBalance(node: tree!) {
//                updateHigh(node: tree!)
//            }else{
//                balance(node: tree!)
//                return
//            }
//            tree = tree!.parent
//        }
//    }
    func updateHigh(node:BinaryTN<T>?) {
        let leftHigh = node?.left?.high ?? 0
        let rightHigh = node?.right?.high ?? 0
        node?.high = max(leftHigh, rightHigh) + 1
    }
    
    func balance(node:BinaryTN<T>){
        let grand = node
        guard let parent = grand.trailChild,let child = parent.trailChild else {
            return
        }
        if parent.isLeft {
            if child.isLeft { // LL
                rotateRight(node: grand)
            }else{ // LR
                rotateLeft(node: parent)
                rotateRight(node: grand)
            }
        }else{ // R
            if child.isRight { // RR
                rotateLeft(node: grand)
            }else{ // RL
                rotateRight(node: parent)
                rotateLeft(node: grand)
            }
        }
    }
    func isBalance(node:BinaryTN<T>) -> Bool {
        let leftHigh = node.left?.high ?? 0
        let rightHigh = node.right?.high ?? 0
        if abs((leftHigh - rightHigh)) > 1 {
            return false
        }
        return true
    }
    
    func rotateRight(node:BinaryTN<T>) {
        let grand = node
        let parent = grand.left
        let child = parent?.right
        
        grand.left = child
        parent?.right = grand
        parent?.parent = grand.parent
        
        if grand.isLeft {
            grand.parent?.left = parent
        }else if grand.isRight {
            grand.parent?.right = parent
        }else{
            root = parent
        }
        grand.parent = parent
        child?.parent = grand
        
        updateHigh(node: child)
        updateHigh(node: grand)
    }
    
    func rotateLeft(node:BinaryTN<T>) {
        let grand = node
        let parent = grand.right
        let child = parent?.left
        
        grand.right = child
        parent?.left = grand
        parent?.parent = grand.parent
        
        if grand.isLeft {
            grand.parent?.left = parent
        }else if grand.isRight {
            grand.parent?.right = parent
        }else{
            root = parent
        }
        grand.parent = parent
        child?.parent = grand
        
        updateHigh(node: child)
        updateHigh(node: grand)
        
    }
    
}


class BinaryTN<T:Equatable> {
    var left:BinaryTN?
    var right:BinaryTN?
    var element:T?
    var high:Int = 0
    var parent:BinaryTN?
    var color:String = "red"
    var asString:String {return treeString(self){("\($0.element!)",$0.left,$0.right)}}
    var asRBString:String {return treeString(self){("\($0.color)_\($0.element!)",$0.left,$0.right)}}

    init(value:T,parent:BinaryTN?) {
        element = value
        self.parent = parent
    }
    var isLeft:Bool{
        get{
            if element == parent?.left?.element {
                return true
            }
            return false
        }
    }
    var isRight:Bool{
        get{
            if element == parent?.right?.element {
                return true
            }
            return false
        }
    }
    var trailChild:BinaryTN?{
        get{
            let leftHigh = left?.high ?? 0
            let rightHigh = right?.high ?? 0
            if leftHigh > rightHigh {
                return left
            }
            if rightHigh > leftHigh {
                return right
            }
            if isLeft {
                return left
            }
            if isRight {
                return right
            }
            return nil
        }
    }
    
    func blackFun(){
        color = "black"
    }
    func redFun(){
        color = "red"
    }
    func isRed() -> Bool {
        return color == "red"
    }
    func isBlack() -> Bool {
        return color == "black"
    }
    
    func sibing() -> BinaryTN<T>? {
        if isLeft {
            return parent?.right
        }
        if isRight  {
            return parent?.left
        }
        return nil
    }
}
